#!/usr/bin/env python

# Created on 2018/12
# Author: Kaituo XU

import argparse


from data import AudioDataLoader, AudioDataset
from solver import Solver
from model import ConvTasNet
from mindspore import Model
from data_test import DatasetGenerator
import mindspore.dataset as ds
from mindspore import nn
from mindspore.train.callback import LossMonitor
from network_define import WithLossCell
from Loss_final1 import loss

parser = argparse.ArgumentParser(
    "Fully-Convolutional Time-domain Audio Separation Network (Conv-TasNet) "
    "with Permutation Invariant Training")
# General config
# Task related
#parser.add_argument('--train_dir', type=str, default='/home/heu_MEDAI/RenQQ/The last/src/out/tr',
                   # help='directory including mix.json, s1.json and s2.json')
parser.add_argument('--train_dir', type=str, default='/mass_data/dataset/LS-2mix/Libri2Mix/tr',
                     help='directory including mix.json, s1.json and s2.json')

parser.add_argument('--valid_dir', type=str, default='/home/heu_MEDAI/RenQQ/The last/src/out/cv',
                    help='directory including mix.json, s1.json and s2.json')
parser.add_argument('--sample_rate', default=8000, type=int,
                    help='Sample rate')
parser.add_argument('--segment', default=4, type=float,   #取音频的长度，2s。#数据集语音长度要相同
                    help='Segment length (seconds)')
parser.add_argument('--cv_maxlen', default=8, type=float,
                    help='max audio length (seconds) in cv, to avoid OOM issue.')  #最大音频长度，防止溢出
# Network architecture
parser.add_argument('--N', default=256, type=int,#256
                    help='Number of filters in autoencoder')
parser.add_argument('--L', default=20, type=int,
                    help='Length of the filters in samples (40=5ms at 8kHZ)')
parser.add_argument('--B', default=256, type=int,#256
                    help='Number of channels in bottleneck 1 × 1-conv block')  #1 × 1-conv的通道数
parser.add_argument('--H', default=512, type=int,#512
                    help='Number of channels in convolutional blocks')   #卷积块的通道数
parser.add_argument('--P', default=3, type=int,
                    help='Kernel size in convolutional blocks')        #卷积核的大小
parser.add_argument('--X', default=8, type=int,
                    help='Number of convolutional blocks in each repeat') #每次8个卷几块
parser.add_argument('--R', default=4, type=int,
                    help='Number of repeats') # 重复4次
parser.add_argument('--C', default=2, type=int,
                    help='Number of speakers') #说话者数量
parser.add_argument('--norm_type', default='gLN', type=str,
                    choices=['gLN', 'cLN', 'BN'], help='Layer norm type')  #归一化的方法
parser.add_argument('--causal', type=int, default=0,
                    help='Causal (1) or noncausal(0) training')  #因果设定
parser.add_argument('--mask_nonlinear', default='relu', type=str,
                    choices=['relu', 'softmax'], help='non-linear to generate mask')  #产生mask的非线性层
# Training config
parser.add_argument('--use_cuda', type=int, default=0,
                    help='Whether use GPU') #是否使用GPU，0==no
parser.add_argument('--epochs', default=10, type=int,
                    help='Number of maximum epochs')
parser.add_argument('--half_lr', dest='half_lr', default=0, type=int,
                    help='Halving learning rate when get small improvement')
parser.add_argument('--early_stop', dest='early_stop', default=0, type=int,
                    help='Early stop training when no improvement for 10 epochs') #自动停止，如果10个epoch没有提升
parser.add_argument('--max_norm', default=5, type=float,
                    help='Gradient norm threshold to clip')
# minibatch
parser.add_argument('--shuffle', default=0, type=int,
                    help='reshuffle the data at every epoch')
parser.add_argument('--batch_size', default=2, type=int,
                    help='Batch size')
parser.add_argument('--num_workers', default=4, type=int,
                    help='Number of workers to generate minibatch')
# optimizer
parser.add_argument('--optimizer', default='adam', type=str,
                    choices=['sgd', 'adam'],
                    help='Optimizer (support sgd and adam now)')
parser.add_argument('--lr', default=1e-3, type=float,
                    help='Init learning rate')
parser.add_argument('--momentum', default=0.0, type=float,
                    help='Momentum for optimizer')
parser.add_argument('--l2', default=0.0, type=float,
                    help='weight decay (L2 penalty)')
# save and load model  #保存模型
parser.add_argument('--save_folder', default='exp/temp',
                    help='Location to save epoch models')
parser.add_argument('--checkpoint', dest='checkpoint', default=1, type=int,
                    help='Enables checkpoint saving of model')
parser.add_argument('--continue_from', default='',
                    help='Continue from checkpoint model')
parser.add_argument('--model_path', default='final.pth.tar',
                    help='Location to save best validation model')
# logging
parser.add_argument('--print_freq', default=10, type=int,
                    help='Frequency of printing training infomation')
parser.add_argument('--visdom', dest='visdom', type=int, default=0,
                    help='Turn on visdom graphing')
parser.add_argument('--visdom_epoch', dest='visdom_epoch', type=int, default=0,
                    help='Turn on visdom graphing each epoch')
parser.add_argument('--visdom_id', default='TasNet training',
                    help='Identifier for visdom run')


def main(args):
    # Construct Solver
    # data
    # tr_dataset = AudioDataset(args.train_dir, args.batch_size,
    #                           sample_rate=args.sample_rate, segment=args.segment)
    # cv_dataset = AudioDataset(args.valid_dir, batch_size=1,  # 1 -> use less GPU memory to do cv
    #                           sample_rate=args.sample_rate,
    #                           segment=-1, cv_maxlen=args.cv_maxlen)  # -1 -> use full audio
    # tr_loader = AudioDataLoader(tr_dataset, batch_size=1,
    #                             shuffle=args.shuffle,
    #                             num_workers=args.num_workers)
    # cv_loader = AudioDataLoader(cv_dataset, batch_size=1,
    #                             num_workers=0)
    tr_dataset = DatasetGenerator(args.train_dir, args.batch_size,
                                  sample_rate=args.sample_rate, segment=args.segment)
    tr_loader = ds.GeneratorDataset(tr_dataset, ["mixture", "lens", "sources"], shuffle=False)
    tr_loader = tr_loader.batch(batch_size=1)

    # cv_dataset = DatasetGenerator(args.valid_dir, args.batch_size,
    #                               sample_rate=args.sample_rate, L=args.L)
    # cv_loader = ds.GeneratorDataset(cv_dataset, ["mixture", "lens", "sources"], shuffle=False)
    # data = {'tr_loader': tr_loader, 'cv_loader': cv_loader}
    # model
    # model = ConvTasNet(args.N, args.L, args.B, args.H, args.P, args.X, args.R,
    #                    args.C, norm_type=args.norm_type, causal=args.causal,
    #                    mask_nonlinear=args.mask_nonlinear)
    net = ConvTasNet(args.N, args.L, args.B, args.H, args.P, args.X, args.R,
                     args.C, norm_type=args.norm_type, causal=args.causal,
                     mask_nonlinear=args.mask_nonlinear)
    #print(net)
    net = net.set_train()
    optimizier = nn.Adam(net.get_parameters(), learning_rate=args.lr, weight_decay=args.l2)
    my_loss = loss()
    cb = LossMonitor()
    net_with_loss = WithLossCell(net, my_loss)
    model = Model(net_with_loss, optimizer=optimizier)
    model.train(epoch=10, train_dataset=tr_loader, callbacks=cb, dataset_sink_mode=False)
    # # optimizer
    # if args.optimizer == 'sgd':
    #     optimizier = torch.optim.SGD(model.parameters(),
    #                                  lr=args.lr,
    #                                  momentum=args.momentum,
    #                                  weight_decay=args.l2)
    # elif args.optimizer == 'adam':
    #     optimizier = torch.optim.Adam(model.parameters(),
    #                                   lr=args.lr,
    #                                   weight_decay=args.l2)
    # else:
    #     print("Not support optimizer")
    #     return
    #
    # # solver
    # solver = Solver(data, model, optimizier, args)
    # solver.train()


if __name__ == '__main__':
    from mindspore import context
    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", pynative_synchronize=True, device_id=6)
    args = parser.parse_args()
    print(args)
    main(args)

